import re
import os
import logging
from typing import Dict, List, Tuple, Optional, Set

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('CppFlowMap')

class CppCodeParser:
    """Improved C/C++ code parser with better main function detection."""
    
    def __init__(self):
        self.functions = {}
        self.entry_points = []
        self.debug_mode = True
        self.logger = logger
        
        # Define patterns for main function detection
        self.main_patterns = [
            # Standard C/C++ main function patterns
            r'int\s+main\s*\(\s*int\s+\w+\s*,\s*char\s*\*+\s*\w+\[\s*\]\s*\)\s*\{',  # int main(int argc, char *argv[]) {
            r'int\s+main\s*\(\s*int\s+\w+\s*,\s*char\s*\*+\s*\*+\s*\w+\s*\)\s*\{',    # int main(int argc, char **argv) {
            r'int\s+main\s*\(\s*void\s*\)\s*\{',                                     # int main(void) {
            r'int\s+main\s*\(\s*\)\s*\{',                                            # int main() {
            r'int\s+main\s*\(\s*int\s+\w+\s*\)\s*\{'                                 # int main(int argc) {
        ]
    
    def find_main_function(self, files) -> Tuple[Optional[str], Optional[str]]:
        """Find the main function in the provided C/C++ source files."""
        self.logger.info(f"Searching for main function in {len(files)} files")
        
        # First try looking for an exact match to a main function
        for file_path in files:
            if not file_path.endswith(('.c', '.cpp', '.cxx', '.cc', '.h', '.hpp')):
                continue
                
            try:
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
                    content = file.read()
                    
                    for pattern in self.main_patterns:
                        matches = re.finditer(pattern, content, re.MULTILINE)
                        for match in matches:
                            # Found a main function
                            self.logger.info(f"Found main function in {file_path}")
                            return "main", file_path
            except Exception as e:
                self.logger.error(f"Error reading file {file_path}: {e}")
        
        # If we get here, no main function was found with the standard patterns
        # Let's try a more thorough search
        self.logger.warning("Standard main function not found. Trying alternative detection methods.")
        
        # Look for files named main.c or main.cpp
        main_files = [f for f in files if os.path.basename(f).lower().startswith('main.')]
        for file_path in main_files:
            try:
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
                    content = file.read()
                    
                    # Try a more lenient pattern that might catch non-standard main functions
                    matches = re.finditer(r'(?:int|void)\s+main\s*\([^)]*\)\s*\{', content, re.MULTILINE)
                    for match in matches:
                        self.logger.info(f"Found main function with alternative pattern in {file_path}")
                        return "main", file_path
                    
                    # Look for lines matching 'main' to help with debugging
                    matches = re.finditer(r'\bmain\b', content, re.MULTILINE)
                    for match in matches:
                        line_start = content.rfind('\n', 0, match.start()) + 1
                        line_end = content.find('\n', match.end())
                        if line_end == -1:
                            line_end = len(content)
                        line = content[line_start:line_end].strip()
                        self.logger.debug(f"Found 'main' reference: {line}")
            except Exception as e:
                self.logger.error(f"Error reading file {file_path}: {e}")
        
        # Fallback to any function with "main" in its name
        for func_name in self.functions:
            if "main" in func_name.lower():
                self.logger.warning(f"Main function not found. Using '{func_name}' as a fallback.")
                return func_name, "unknown"
                
        # If no main function is found at all
        if self.functions:
            first_func = next(iter(self.functions))
            self.logger.warning(f"Main function not found. Using the first found function '{first_func}' as entry point.")
            return first_func, "unknown"
            
        return None, None

    def extract_main_from_file(self, file_path: str) -> Optional[str]:
        """Extract the main function definition from a specific file."""
        if not os.path.exists(file_path):
            self.logger.error(f"File not found: {file_path}")
            return None
            
        try:
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
                content = file.read()
                
                # Try to find the main function
                for pattern in self.main_patterns:
                    match = re.search(pattern, content, re.MULTILINE)
                    if match:
                        # Find the full function body (accounting for nested braces)
                        start_pos = match.start()
                        pos = match.end()
                        brace_count = 1
                        
                        while pos < len(content) and brace_count > 0:
                            if content[pos] == '{':
                                brace_count += 1
                            elif content[pos] == '}':
                                brace_count -= 1
                            pos += 1
                        
                        if brace_count == 0:
                            main_func = content[start_pos:pos]
                            self.logger.info(f"Successfully extracted main function ({len(main_func)} chars)")
                            return main_func
        except Exception as e:
            self.logger.error(f"Error extracting main function from {file_path}: {e}")
            
        return None
        
    def dump_file_overview(self, file_path: str):
        """Dump an overview of a file to help with debugging."""
        if not os.path.exists(file_path):
            self.logger.error(f"File not found: {file_path}")
            return
            
        try:
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
                content = file.read()
                lines = content.split('\n')
                
                self.logger.info(f"File overview for {file_path}: {len(lines)} lines, {len(content)} chars")
                
                # Look for potential function definitions
                func_patterns = [
                    r'\b(\w+)\s+(\w+)\s*\([^)]*\)\s*\{',  # return_type name(...) {
                    r'\b(\w+)::\s*(\w+)\s*\([^)]*\)\s*\{' # class::name(...) {
                ]
                
                for pattern in func_patterns:
                    matches = re.finditer(pattern, content, re.MULTILINE)
                    for match in matches:
                        line_num = content[:match.start()].count('\n') + 1
                        self.logger.info(f"Line {line_num}: Found potential function: {match.group(0)[:50]}...")
        except Exception as e:
            self.logger.error(f"Error analyzing file {file_path}: {e}")

# Helper function to analyze a directory of C/C++ files
def analyze_cpp_directory(directory: str) -> CppCodeParser:
    """Analyze a directory of C/C++ files to find the main function."""
    parser = CppCodeParser()
    
    # Find all C/C++ files in the directory
    cpp_files = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith(('.c', '.cpp', '.cxx', '.cc', '.h', '.hpp')):
                cpp_files.append(os.path.join(root, file))
    
    # Try to find the main function
    main_func, main_file = parser.find_main_function(cpp_files)
    
    if main_func:
        logger.info(f"Found main function: {main_func} in {main_file}")
        
        # If we found a main function by name but not the actual code, try to extract it
        if main_file != "unknown":
            main_code = parser.extract_main_from_file(main_file)
            if main_code:
                logger.info(f"Extracted main function code: {len(main_code)} chars")
    else:
        logger.warning("No main function found")
        
        # If no main function was found, dump info about some key files to help debugging
        for file in cpp_files:
            if os.path.basename(file).lower().startswith('main.'):
                logger.info(f"Analyzing potential main file: {file}")
                parser.dump_file_overview(file)
    
    return parser

    def parse_all_files(self, source_dirs, target_file=None):
        # ... existing code ...
        
        # After parsing all files, find the main function
        main_function, main_file = self.find_main_function(all_files)
        if main_function:
            self.root_functions = [main_function]
            self.logger.info(f"Using '{main_function}' from {main_file} as the main entry point.")
        else:
            self.logger.error("No functions found in the codebase. Unable to build call tree.")
        
        # ... rest of the method ... 